'''
Author: Yan Jianhang
Date: 2022-04-12 19:51:21
LastEditors: Yan Jianhang
LastEditTime: 2022-04-29 00:01:09
FilePath: /STransGrasp/lib/tf/vis_points.py
Description: 可视化 grasp transfer 后的抓取点。(grasp_data/*/grasp_points*.npz)
'''

import pickle, argparse, os, yaml, sys, torch
import numpy as np
import open3d as o3d
import open3d.visualization as vis
sys.path.append('./')
from lib.fk.FK_layer import FK_layer
from lib.fk.utils import o3d_read_mesh
from lib.tf.utils import get_grasp_hand_lineset

parser = argparse.ArgumentParser()
parser.add_argument('--data_root', type=str, default='datasets')
parser.add_argument('--category', type=str, default='mug')
parser.add_argument('--mode', type=str, default='eval')
parser.add_argument('--grasp_id', type=str, default='155')
opt = parser.parse_args()

def label_grasp():
    with open('cfg/class_cfg.yml', 'r') as stream:
        cls_cfg = yaml.safe_load(stream)
    inst_name = cls_cfg[opt.category]['src_inst_name']
    inst_path = f"{opt.category}/train/{inst_name}/0"
    grasp_file = f"grasp_data/{inst_path}/label_grasp_{opt.grasp_id}.npz"
    grasp_pose = np.load(grasp_file)
    palm_q = grasp_pose['palm_q'][[3,0,1,2]]       # [x,y,z,w] -> [w,x,y,z]
    base = np.concatenate([grasp_pose['palm_t'], palm_q], axis=0)
    rotations = grasp_pose['joints'][2:24]
    rotations = torch.FloatTensor(rotations).reshape(1, -1).cuda()
    base = torch.FloatTensor(base).reshape(1, -1).cuda()    
    model = FK_layer(base, rotations)
    model.to('cuda')
    positions, transformed_pts = model()
    pcd = o3d.geometry.PointCloud()
    transformed_pts = transformed_pts[0].detach().cpu().numpy()
    transformed_pts_offset = transformed_pts + np.array([0, 0.07, 0])
    pcd.points = o3d.utility.Vector3dVector(transformed_pts_offset)
    
    contact_info_path = f"grasp_data/{inst_path}/label_grasp_{opt.grasp_id}_contact_info.npz"
    contact_label = np.load(contact_info_path)
    
    positions =  positions.detach().cpu().numpy()
    # hand_skeleton = get_grasp_hand_lineset(positions)
    
    return pcd, transformed_pts_offset, contact_label

if __name__ == "__main__":
    
    hand_pc_o3d, hand_pc_np, contact_label = label_grasp()
    
    cnt = 0
    vis_file_name = f"grasp_{opt.grasp_id}_points_on_surface"  # "grasp_points": free points, "grasp_points_on_surface": surface points
    for inst in os.listdir("grasp_data/{}/{}".format(opt.category, opt.mode)):
        for part in os.listdir("grasp_data/{}/{}/{}".format(opt.category, opt.mode, inst)):
            if cnt != 51:
                cnt+=1
                continue
            
            inst_path = f"{opt.category}/{opt.mode}/{inst}/{part}"
            
            show_list = [hand_pc_o3d, ]
            g_pts_file = f"grasp_data/{inst_path}/{vis_file_name}.npz"
            g_pts = np.load(g_pts_file)
            
            scale = g_pts["obj_scale"]
            
            # show object mesh
            inst_path = inst_path.replace('eval_pc', 'eval')
            obj_file   = f"datasets/obj/{inst_path}.ply"
            mesh = o3d_read_mesh(obj_file, scale=scale)
            show_list.append(mesh)

            # touched set on target object
            grasp_points = o3d.geometry.PointCloud()
            g_pts_scaled = g_pts["grasp_points"] * scale * 0.5
            grasp_points.points = o3d.utility.Vector3dVector(g_pts_scaled)
            grasp_points.paint_uniform_color([1,0,0])
            show_list.append(grasp_points)

            # transfer goal, shown by line set
            cor_line = o3d.geometry.LineSet()
            contact_pts_num = g_pts_scaled.shape[0]
            assert contact_pts_num == hand_pc_np[contact_label['contact_idx']].shape[0]
            pts = np.concatenate([hand_pc_np[contact_label['contact_idx']], g_pts_scaled], axis=0)
            cor_line.points = o3d.utility.Vector3dVector(pts)
            line_np = np.zeros([contact_pts_num, 2], dtype=np.int)
            for i in range(0, contact_pts_num):
                line_np[i][0] = i
                line_np[i][1] = i + contact_pts_num
            cor_line.lines = o3d.utility.Vector2iVector(line_np)
            show_list.append(cor_line)
            
            recon_file = f"dif/eval/{opt.category}/{inst}/0/checkpoints/test.ply"
            if os.path.exists(recon_file):
                recon_mesh = o3d_read_mesh(recon_file, scale=scale * 0.5)
                recon_mesh.paint_uniform_color([0,0,0.5])
                show_list.append(recon_mesh)

            vis.draw(show_list)